/* 
just run:
g++ -std=c++17 -pthread -Wall -O2 core2.cpp && ./a.out 
*/

#include <iostream>
#include <string>
#include <vector>
#include <memory>
#include <mutex>
#include <shared_mutex>
#include <filesystem>
#include <fstream>
#include <cstring>
#include <sys/mman.h>
#include <fcntl.h>
#include <unistd.h>
#include <system_error>
#include <functional>
#include <map>
#include <set>

namespace banker {

constexpr size_t CHUNK_SIZE = 4 * 1024 * 1024; // 4MB
constexpr size_t NUM_CHUNKS = 256; // 256 chunks = 1GB

// 异常类，用于包装系统错误
class MemoryMapException : public std::runtime_error {
public:
    MemoryMapException(const std::string& msg, int errno_val)
        : std::runtime_error(msg + ": " + std::strerror(errno_val)), error_code(errno_val) {}
    
    int code() const { return error_code; }
    
private:
    int error_code;
};

// 管理延迟映射的内存区域
class LazyMappedMemory {
public:
    explicit LazyMappedMemory(const std::string& base_path) : base_path(base_path) {
        // 确保基础目录存在
        std::filesystem::create_directories(base_path);
        
        // 初始化映射和文件描述符
        mappings.resize(NUM_CHUNKS, nullptr);
        file_descriptors.resize(NUM_CHUNKS, -1);
    }
    
    ~LazyMappedMemory() {
        // 清理所有映射和文件描述符
        for (size_t i = 0; i < NUM_CHUNKS; ++i) {
            if (mappings[i] != nullptr) {
                munmap(mappings[i], CHUNK_SIZE);
            }
            if (file_descriptors[i] != -1) {
                close(file_descriptors[i]);
            }
        }
    }
    
    // 禁用拷贝和移动
    LazyMappedMemory(const LazyMappedMemory&) = delete;
    LazyMappedMemory& operator=(const LazyMappedMemory&) = delete;
    LazyMappedMemory(LazyMappedMemory&&) = delete;
    LazyMappedMemory& operator=(LazyMappedMemory&&) = delete;
    
    uint8_t get_byte(size_t offset) {
        if (offset >= CHUNK_SIZE * NUM_CHUNKS) {
            return 0;
        }
        
        const size_t chunk_idx = offset / CHUNK_SIZE;
        const size_t chunk_offset = offset % CHUNK_SIZE;
        
        std::shared_lock lock(mutex);
        
        // 如果映射不存在，尝试创建或返回0
        if (mappings[chunk_idx] == nullptr) {
            lock.unlock(); // 释放共享锁，获取独占锁
            std::unique_lock unique_lock(mutex);
            
            // 双重检查，防止竞态条件
            if (mappings[chunk_idx] == nullptr) {
                try {
                    ensure_mapping(chunk_idx, false);
                } catch (const MemoryMapException&) {
                    // 文件不存在，返回0
                    return 0;
                }
            }
            unique_lock.unlock();
            lock.lock(); // 重新获取共享锁
        }
        
        return static_cast<uint8_t*>(mappings[chunk_idx])[chunk_offset];
    }
    
    void set_byte(size_t offset, uint8_t value) {
        if (offset >= CHUNK_SIZE * NUM_CHUNKS) {
            return;
        }
        
        const size_t chunk_idx = offset / CHUNK_SIZE;
        const size_t chunk_offset = offset % CHUNK_SIZE;
        
        std::unique_lock lock(mutex);
        
        // 确保映射存在（如果需要则创建文件）
        if (mappings[chunk_idx] == nullptr) {
            ensure_mapping(chunk_idx, true);
        }
        
        // 写入数据
        static_cast<uint8_t*>(mappings[chunk_idx])[chunk_offset] = value;
        
        // 标记块为脏
        dirty_chunks.insert(chunk_idx);
    }
    
    // 读取64位值
    int64_t get_int64(size_t offset) {
        int64_t value = 0;
        for (int i = 0; i < 8; ++i) {
            value |= static_cast<int64_t>(get_byte(offset + i)) << (i * 8);
        }
        return value;
    }
    
    // 写入64位值
    void set_int64(size_t offset, int64_t value) {
        for (int i = 0; i < 8; ++i) {
            set_byte(offset + i, (value >> (i * 8)) & 0xFF);
        }
    }
    
    // 将所有修改刷新到磁盘
    void flush_all() {
        std::shared_lock lock(mutex);
        for (size_t i = 0; i < NUM_CHUNKS; ++i) {
            if (mappings[i] != nullptr) {
                msync(mappings[i], CHUNK_SIZE, MS_SYNC);
            }
        }
        dirty_chunks.clear();
    }
    
    // 刷新所有脏块到磁盘
    void flush_dirty_chunks() {
        std::unique_lock lock(mutex);
        for (auto chunk_idx : dirty_chunks) {
            if (mappings[chunk_idx] != nullptr) {
                msync(mappings[chunk_idx], CHUNK_SIZE, MS_SYNC);
            }
        }
        dirty_chunks.clear();
    }
    
    // 将指定块刷新到磁盘
    void flush_chunk(size_t chunk_idx) {
        if (chunk_idx >= NUM_CHUNKS) {
            return;
        }
        
        std::shared_lock lock(mutex);
        if (mappings[chunk_idx] != nullptr) {
            msync(mappings[chunk_idx], CHUNK_SIZE, MS_SYNC);
        }
        dirty_chunks.erase(chunk_idx);
    }
    
    // 卸载指定块（释放内存映射）
    void unload_chunk(size_t chunk_idx) {
        if (chunk_idx >= NUM_CHUNKS) {
            return;
        }
        
        std::unique_lock lock(mutex);
        if (mappings[chunk_idx] != nullptr) {
            munmap(mappings[chunk_idx], CHUNK_SIZE);
            mappings[chunk_idx] = nullptr;
        }
        dirty_chunks.erase(chunk_idx);
        // 注意：我们保留文件描述符以便后续快速重新映射
    }
    
private:
    void ensure_mapping(size_t chunk_idx, bool create_if_missing) {
        if (chunk_idx >= NUM_CHUNKS) {
            throw std::out_of_range("Chunk index out of range");
        }
        
        const std::string file_path = base_path + "/chunk_" + 
                                    (chunk_idx < 10 ? "00" : 
                                     chunk_idx < 100 ? "0" : "") + 
                                    std::to_string(chunk_idx) + ".bin";
        
        // 检查文件是否存在
        if (!std::filesystem::exists(file_path)) {
            if (!create_if_missing) {
                throw MemoryMapException("File does not exist", ENOENT);
            }
            
            // 创建文件并初始化为全零
            int fd = open(file_path.c_str(), O_RDWR | O_CREAT, 0644);
            if (fd == -1) {
                throw MemoryMapException("Failed to create file", errno);
            }
            
            // 设置文件大小为 CHUNK_SIZE
            if (ftruncate(fd, CHUNK_SIZE) == -1) {
                close(fd);
                throw MemoryMapException("Failed to truncate file", errno);
            }
            
            // 初始化文件内容为全零
            std::vector<uint8_t> zeros(CHUNK_SIZE, 0);
            if (write(fd, zeros.data(), zeros.size()) == -1) {
                close(fd);
                throw MemoryMapException("Failed to initialize file", errno);
            }
            
            file_descriptors[chunk_idx] = fd;
        } else if (file_descriptors[chunk_idx] == -1) {
            // 文件存在但尚未打开
            int fd = open(file_path.c_str(), O_RDWR);
            if (fd == -1) {
                throw MemoryMapException("Failed to open file", errno);
            }
            file_descriptors[chunk_idx] = fd;
        }
        
        // 创建内存映射
        void* mapping = mmap(nullptr, CHUNK_SIZE, PROT_READ | PROT_WRITE, 
                            MAP_SHARED, file_descriptors[chunk_idx], 0);
        if (mapping == MAP_FAILED) {
            throw MemoryMapException("Failed to map memory", errno);
        }
        
        mappings[chunk_idx] = mapping;
    }
    
    std::string base_path;
    std::vector<void*> mappings;
    std::vector<int> file_descriptors;
    std::set<size_t> dirty_chunks; // 跟踪哪些块被修改了
    mutable std::shared_mutex mutex;
};

// 提供线程安全版本的包装器
class ThreadSafeLazyMappedMemory {
public:
    explicit ThreadSafeLazyMappedMemory(const std::string& base_path)
        : impl(std::make_shared<LazyMappedMemory>(base_path)) {}
    
    uint8_t get_byte(size_t offset) {
        return impl->get_byte(offset);
    }
    
    void set_byte(size_t offset, uint8_t value) {
        impl->set_byte(offset, value);
    }
    
    int64_t get_int64(size_t offset) {
        return impl->get_int64(offset);
    }
    
    void set_int64(size_t offset, int64_t value) {
        impl->set_int64(offset, value);
    }
    
    void flush_all() {
        impl->flush_all();
    }
    
    void flush_dirty_chunks() {
        impl->flush_dirty_chunks();
    }
    
private:
    std::shared_ptr<LazyMappedMemory> impl;
};

} // namespace banker

// 虚拟机实现
class VirtualMachine {
private:
    // 常量定义
    static const int MEM_LEN = 512;
    static const int CODE_SEG_BEGIN = 0;
    static const int CODE_SEG_LEN = 256;
    static const int DATA_SEG_BEGIN = 256;
    static const int DATA_SEG_LEN = 256;
    static const int COINS_REMAIN_INIT = 2618;
    
    // 全局状态
    banker::ThreadSafeLazyMappedMemory memory;
    int PC = 0;
    std::vector<int64_t> REGSTACK;
    int CYCLES = 0;
    int COINS_REMAIN = COINS_REMAIN_INIT;
    
    // 指令定义
    struct Instruction {
        std::string name;
        int code;
        int args;
        std::function<void()> action;
    };
    
    std::vector<Instruction> ISA;
    std::map<std::string, std::string> ALIAS_MAP;
    
public:
    VirtualMachine(const std::string& base_path) : memory(base_path) {
        initialize_isa();
        initialize_alias_map();
    }
    
    void run() {
        // 加载测试代码
        load_test_code();
        
        // 运行机器
        while (COINS_REMAIN > 0 && PC >= 0 && PC < MEM_LEN) {
            if (!step_machine()) {
                break; // 遇到HALT指令
            }
            CYCLES++;
            
            if (PC >= MEM_LEN) {
                PC = 0;
            }
        }
        
        if (COINS_REMAIN <= 0) {
            std::cout << "NOT-ENOUGH-COINS" << std::endl;
        }
        dump_machine();
    }
    
private:
    void initialize_isa() {
        ISA = {
            {"NOP", 0, 0, [this]() { PC += 1; }},
            {"PUSH", 1, 1, [this]() { 
                int arg = read_code(PC + 1);
                REGSTACK.push_back(arg);
                PC += 2;
            }},
            {"POP", 2, 0, [this]() { 
                if (!REGSTACK.empty()) {
                    REGSTACK.pop_back(); 
                }
                PC += 1;
            }},
            {"ADD", 11, 0, [this]() { 
                if (REGSTACK.size() < 2) return;
                int64_t op2 = REGSTACK.back(); REGSTACK.pop_back();
                int64_t op1 = REGSTACK.back(); REGSTACK.pop_back();
                REGSTACK.push_back(op1 + op2);
                PC += 1;
            }},
            {"SUB", 12, 0, [this]() { 
                if (REGSTACK.size() < 2) return;
                int64_t op2 = REGSTACK.back(); REGSTACK.pop_back();
                int64_t op1 = REGSTACK.back(); REGSTACK.pop_back();
                REGSTACK.push_back(op1 - op2);
                PC += 1;
            }},
            {"MUL", 13, 0, [this]() { 
                if (REGSTACK.size() < 2) return;
                int64_t op2 = REGSTACK.back(); REGSTACK.pop_back();
                int64_t op1 = REGSTACK.back(); REGSTACK.pop_back();
                REGSTACK.push_back(op1 * op2);
                PC += 1;
            }},
            {"DIV", 14, 0, [this]() { 
                if (REGSTACK.size() < 2) return;
                int64_t op2 = REGSTACK.back(); REGSTACK.pop_back();
                int64_t op1 = REGSTACK.back(); REGSTACK.pop_back();
                if (op2 != 0) {
                    REGSTACK.push_back(op1 / op2);
                } else {
                    REGSTACK.push_back(0);
                }
                PC += 1;
            }},
            {"MOD", 15, 0, [this]() { 
                if (REGSTACK.size() < 2) return;
                int64_t op2 = REGSTACK.back(); REGSTACK.pop_back();
                int64_t op1 = REGSTACK.back(); REGSTACK.pop_back();
                if (op2 != 0) {
                    REGSTACK.push_back(op1 % op2);
                } else {
                    REGSTACK.push_back(0);
                }
                PC += 1;
            }},
            {"AND", 21, 0, [this]() { 
                if (REGSTACK.size() < 2) return;
                int64_t op2 = REGSTACK.back(); REGSTACK.pop_back();
                int64_t op1 = REGSTACK.back(); REGSTACK.pop_back();
                REGSTACK.push_back(op1 & op2);
                PC += 1;
            }},
            {"OR", 22, 0, [this]() { 
                if (REGSTACK.size() < 2) return;
                int64_t op2 = REGSTACK.back(); REGSTACK.pop_back();
                int64_t op1 = REGSTACK.back(); REGSTACK.pop_back();
                REGSTACK.push_back(op1 | op2);
                PC += 1;
            }},
            {"NOT", 23, 0, [this]() { 
                if (REGSTACK.empty()) return;
                int64_t op1 = REGSTACK.back(); REGSTACK.pop_back();
                REGSTACK.push_back(~op1);
                PC += 1;
            }},
            {"EQ", 31, 0, [this]() { 
                if (REGSTACK.size() < 2) return;
                int64_t op2 = REGSTACK.back(); REGSTACK.pop_back();
                int64_t op1 = REGSTACK.back(); REGSTACK.pop_back();
                REGSTACK.push_back(op1 == op2 ? 1 : 0);
                PC += 1;
            }},
            {"GT", 32, 0, [this]() { 
                if (REGSTACK.size() < 2) return;
                int64_t op2 = REGSTACK.back(); REGSTACK.pop_back();
                int64_t op1 = REGSTACK.back(); REGSTACK.pop_back();
                REGSTACK.push_back(op1 > op2 ? 1 : 0);
                PC += 1;
            }},
            {"LT", 33, 0, [this]() { 
                if (REGSTACK.size() < 2) return;
                int64_t op2 = REGSTACK.back(); REGSTACK.pop_back();
                int64_t op1 = REGSTACK.back(); REGSTACK.pop_back();
                REGSTACK.push_back(op1 < op2 ? 1 : 0);
                PC += 1;
            }},
            {"LOAD", 41, 1, [this]() { 
                int addr = read_code(PC + 1);
                REGSTACK.push_back(get_data(addr));
                PC += 2;
            }},
            {"STORE", 42, 1, [this]() { 
                int addr = read_code(PC + 1);
                if (!REGSTACK.empty()) {
                    set_data(addr, REGSTACK.back());
                    REGSTACK.pop_back();
                }
                PC += 2;
            }},
            {"LOAD_INDIRECT", 43, 0, [this]() { 
                if (REGSTACK.empty()) return;
                int addr = REGSTACK.back(); REGSTACK.pop_back();
                REGSTACK.push_back(get_data(addr));
                PC += 1;
            }},
            {"STORE_INDIRECT", 44, 0, [this]() { 
                if (REGSTACK.size() < 2) return;
                int addr = REGSTACK.back(); REGSTACK.pop_back();
                int64_t val = REGSTACK.back(); REGSTACK.pop_back();
                set_data(addr, val);
                PC += 1;
            }},
            {"JUMP", 51, 1, [this]() { 
                PC = read_code(PC + 1);
            }},
            {"JUMPIF", 52, 1, [this]() { 
                if (REGSTACK.empty()) {
                    PC += 2;
                } else {
                    int64_t condition = REGSTACK.back(); REGSTACK.pop_back();
                    if (condition == 0) {
                        PC += 2;
                    } else {
                        PC = read_code(PC + 1);
                    }
                }
            }},
            {"PRINT", 61, 0, [this]() { 
                if (!REGSTACK.empty()) {
                    std::cout << REGSTACK.back() << std::endl;
                    REGSTACK.pop_back();
                }
                PC += 1;
            }},
            {"HALT", 255, 0, [this]() {}}
        };
    }
    
    void initialize_alias_map() {
        ALIAS_MAP = {
            {"+", "ADD"}, {"-", "SUB"}, {"*", "MUL"}, {"/", "DIV"}, {"%", "MOD"},
            {"=", "EQ"}, {">", "GT"}, {"<", "LT"},
            {"&", "AND"}, {"|", "OR"}, {"~", "NOT"},
            {"JMP", "JUMP"}, {"HLT", "HALT"}
        };
    }
    
    // 基本设施
    int64_t read_code(int addr) {
        return memory.get_int64((CODE_SEG_BEGIN + addr) * 8);
    }
    
    int64_t get_data(int addr) {
        return memory.get_int64((DATA_SEG_BEGIN + addr) * 8);
    }
    
    void set_data(int addr, int64_t val) {
        memory.set_int64((DATA_SEG_BEGIN + addr) * 8, val);
    }
    
    int64_t get_arg() {
        return read_code(PC + 1);
    }
    
    int64_t get_instr() {
        return read_code(PC);
    }
    
    // 加载测试代码
    void load_test_code() {
        std::vector<std::string> test_code = {
            "PUSH", "100",
            "STORE", "100",
            "LOOP:",
            "LOAD", "100",
            "PUSH", "1",
            "SUB",
            "STORE", "100",
            "LOAD", "100",
            "LOAD", "100",
            "STORE_INDIRECT",
            "LOAD", "100",
            "PUSH", "0",
            "EQ",
            "JUMPIF", "PRIME_LEFT",
            "JUMP", "LOOP",
            "PRIME_LEFT:",
            "PUSH", "2",
            "STORE", "1",
            "P_LOOP:",
            "LOAD", "1",
            "LOAD", "1",
            "MUL",
            "PUSH", "100",
            "GT",
            "JUMPIF", "HALT",
            "LOAD", "1",
            "LOAD_INDIRECT",
            "PUSH", "0",
            "EQ",
            "JUMPIF", "SKIP_MARK",
            "JUMP", "MARK_MULTIPLES",
            "MARK_MULTIPLES:",
            "LOAD", "1",
            "LOAD", "1",
            "MUL",
            "STORE", "0",
            "I_LOOP:",
            "LOAD", "0",
            "PUSH", "100",
            "LT",
            "JUMPIF", "I_CONTINUE",
            "JUMP", "SKIP_MARK",
            "I_CONTINUE:",
            "PUSH", "0",
            "LOAD", "0",
            "STORE_INDIRECT",
            "LOAD", "0",
            "LOAD", "1",
            "ADD",
            "STORE", "0",
            "JUMP", "I_LOOP",
            "SKIP_MARK:",
            "LOAD", "1",
            "PUSH", "1",
            "ADD",
            "STORE", "1",
            "JUMP", "P_LOOP",
            "HALT:",
            "HALT"
        };
        
        // 汇编代码并加载到内存
        auto bytecode = assemble(test_code);
        for (size_t i = 0; i < bytecode.size() && i < CODE_SEG_LEN; i++) {
            memory.set_int64((CODE_SEG_BEGIN + i) * 8, bytecode[i]);
        }
    }
    
    // 简单的汇编器
    std::vector<int64_t> assemble(const std::vector<std::string>& tokens) {
        // 简化的汇编器实现
        // 在实际实现中，您需要处理标签、别名等
        std::vector<int64_t> result;
        std::map<std::string, int> labels;
        
        // 第一遍：收集标签
        int addr = 0;
        for (size_t i = 0; i < tokens.size(); i++) {
            const std::string& token = tokens[i];
            if (token.back() == ':') {
                labels[token.substr(0, token.size() - 1)] = addr;
            } else {
                // 查找指令或别名
                std::string actual_name = token;
                if (ALIAS_MAP.find(token) != ALIAS_MAP.end()) {
                    actual_name = ALIAS_MAP[token];
                }
                
                // 查找指令
                for (const auto& instr : ISA) {
                    if (instr.name == actual_name) {
                        addr++;
                        if (instr.args > 0) {
                            i++;
                            addr++;
                        }
                        break;
                    }
                }
            }
        }
        
        // 第二遍：生成代码
        for (size_t i = 0; i < tokens.size(); i++) {
            const std::string& token = tokens[i];
            if (token.back() == ':') {
                continue; // 跳过标签定义
            }
            
            // 查找指令或别名
            std::string actual_name = token;
            if (ALIAS_MAP.find(token) != ALIAS_MAP.end()) {
                actual_name = ALIAS_MAP[token];
            }
            
            // 查找指令
            for (const auto& instr : ISA) {
                if (instr.name == actual_name) {
                    result.push_back(instr.code);
                    
                    if (instr.args > 0) {
                        i++;
                        if (i < tokens.size()) {
                            const std::string& arg = tokens[i];
                            if (labels.find(arg) != labels.end()) {
                                result.push_back(labels[arg]);
                            } else {
                                try {
                                    result.push_back(std::stoi(arg));
                                } catch (...) {
                                    result.push_back(0);
                                }
                            }
                        } else {
                            result.push_back(0);
                        }
                    }
                    break;
                }
            }
        }
        
        return result;
    }
    
    // 调试函数
    void dump_machine() {
        std::cout << "((CODE-SEG-BEGIN " << CODE_SEG_BEGIN << ")" << std::endl;
        std::cout << " (CODE-SEG-LEN " << CODE_SEG_LEN << ")" << std::endl;
        std::cout << " (DATA-SEG-BEGIN " << DATA_SEG_BEGIN << ")" << std::endl;
        std::cout << " (DATA-SEG-LEN " << DATA_SEG_LEN << ")" << std::endl;
        std::cout << " (PC " << PC << ")" << std::endl;
        std::cout << " (REGSTACK (";
        for (size_t i = 0; i < REGSTACK.size(); i++) {
            if (i > 0) std::cout << " ";
            std::cout << REGSTACK[i];
        }
        std::cout << "))" << std::endl;
        std::cout << " (COINS-REMAIN " << COINS_REMAIN << ")" << std::endl;
        std::cout << " (CYCLES " << CYCLES << "))" << std::endl;
    }
    
    void debug_print() {
        std::map<int, std::string> dasm_map;
        for (const auto& instr : ISA) {
            dasm_map[instr.code] = instr.name;
        }
        
        int instr_code = get_instr();
        std::string instr_name = dasm_map.count(instr_code) ? dasm_map[instr_code] : "UNKNOWN";
        
        std::cout << "(CYCLES: " << CYCLES << " REGSTACK: (";
        for (size_t i = 0; i < REGSTACK.size(); i++) {
            if (i > 0) std::cout << " ";
            std::cout << REGSTACK[i];
        }
        std::cout << ") instr: " << instr_name;
        
        for (const auto& instr : ISA) {
            if (instr.code == instr_code && instr.args > 0) {
                std::cout << " " << get_arg();
                break;
            }
        }
        std::cout << ")" << std::endl;
    }
    
    bool step_machine() {
        debug_print();
        
        int instr_code = get_instr();
        
        // 查找并执行指令
        for (const auto& instr : ISA) {
            if (instr.code == instr_code) {
                // 执行指令
                instr.action();
                
                // 确保指令的原子性：刷新所有脏块到磁盘
                memory.flush_dirty_chunks();
                
                COINS_REMAIN--;
                return instr_code != 255; // 如果不是HALT指令，返回true
            }
        }
        
        // 未知指令，跳过
        PC += 1;
        COINS_REMAIN--;
        return true;
    }
};

int main() {
    try {
        const std::string base_path = "./vm_memory";
        
        // 创建虚拟机实例
        VirtualMachine vm(base_path);
        
        // 运行虚拟机
        vm.run();
        
    } catch (const std::exception& e) {
        std::cerr << "Error: " << e.what() << std::endl;
        return 1;
    }
    
    return 0;
}
